# 20250211 一些用于2D问题添加的函数
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as tri
from matplotlib.patches import ConnectionPatch

import torch

def add_axes(ax, x, y_pred, y, xlim0, xlim1, ylim0, ylim1):
    axins = ax.inset_axes((0.2, 0.2, 0.4, 0.3))
    # xlim0, xlim1 = 1.2, 1.3
    # ylim0, ylim1 = -1, 1
    axins.set_xlim(xlim0, xlim1)  # 设置x轴的范围
    axins.set_ylim(ylim0, ylim1)  # 设置y轴的范围
    axins.plot(x, y, "k-")
    axins.plot(x, y_pred, "r--")

    tx0 = xlim0
    tx1 = xlim1
    ty0 = ylim0
    ty1 = ylim1
    sx = [tx0,tx1,tx1,tx0,tx0]
    sy = [ty0,ty0,ty1,ty1,ty0]
    ax.plot(sx,sy,"black")

    # 画两条线
    xy = (xlim0,ylim0)
    xy2 = (xlim0,ylim1)
    con = ConnectionPatch(xyA=xy2,xyB=xy,coordsA="data",coordsB="data",
            axesA=axins,axesB=ax)
    axins.add_artist(con)

    xy = (xlim1,ylim0)
    xy2 = (xlim1,ylim1)
    con = ConnectionPatch(xyA=xy2,xyB=xy,coordsA="data",coordsB="data",
            axesA=axins,axesB=ax)
    axins.add_artist(con)
    return
    
def plot_tri_contourf(x1, x2, value, title='title', shape=False, equal_ratio=True):

    if type(x1) is torch.Tensor:
        x1 = x1.numpy()
        x2 = x2.numpy()
        if value.requires_grad:
            value  = value.detach()
        value = value.cpu().numpy()
    
    # Create the Triangulation; no triangles so Delaunay triangulation created.
    triangulation = tri.Triangulation(x1, x2)

    # Mask off unwanted triangles.
    if shape == 'annulus':
        tol = 1e-7 + 1.0
        if np.max(np.abs(x1))>tol or np.max(np.abs(x2)) > tol:
            raise ValueError('Only support [-1,1] range for annulus.')

        # Mask off unwanted triangles.
        min_radius = 0.5
        triangulation.set_mask(np.hypot(x1[triangulation.triangles].mean(axis=1),
                                        x2[triangulation.triangles].mean(axis=1)) < min_radius)

    if shape == 'L-shape':
        tol = 1e-7 + 1.0
        if np.max(np.abs(x1)) > tol or np.max(np.abs(x2)) > tol:
            raise ValueError('Only support [-1,1] range for annulus.')

        dx = x2[1] - x2[0]
        v1 = x1[triangulation.triangles]
        v2 = x2[triangulation.triangles]

        x1_diff = np.stack([v1[:, 0] - v1[:, 1], v1[:, 1] - v1[:, 2], v1[:, 2] - v1[:, 0]], axis=1)
        x2_diff = np.stack([v2[:, 0] - v2[:, 1], v2[:, 1] - v2[:, 2], v2[:, 2] - v2[:, 0]], axis=1)

        tol = dx * 1.2
        exceed_ct = np.sum((x1_diff > tol) * 1 + (x2_diff > tol) * 1, axis=1)
        bd_mask = exceed_ct > 0
        triangulation.set_mask(bd_mask)

    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111)
    if equal_ratio:
         ax.set_aspect('equal')

    tpc = ax.tricontourf(triangulation, value, levels=20, cmap='jet')   # ‘rainbow’, 'jet'
    fig.colorbar(tpc)
    ax.set_title(title)
    ax.set_xlabel('x1')
    ax.set_ylabel('x2')
    ax.grid(which='both')

    return fig

def plot_2d_mesh(x1, x2, value, title='title', figsize=(8,6), equal_ratio=True, z_min=None, z_max=None):
    #  plot
    fig, (ax, cax) = plt.subplots(ncols=2,figsize=figsize, gridspec_kw={"width_ratios": [1, 0.05]})
    fig.subplots_adjust(wspace=0.1)
    
    im = ax.pcolormesh(x1, x2, value, cmap='jet', shading='auto', vmin=z_min, vmax=z_max)

    # ax.plot(0,0,'r*')
    fig.colorbar(im, cax=cax)
    if equal_ratio:
        ax.set_aspect('equal')
    ax.set_title(title)
    ax.grid(True)
    plt.show()


class Domain():
    def __init__(self, domain_range, domain_shape='box', mesh_domain=None, mesh_boundary=None):
        # domain_range ex: (-1,1,-1,1)
        domain_range = np.array(domain_range)
        domain_range = domain_range.reshape((domain_range.shape[0] // 2, 2))
        self.domain_range = domain_range
        self.domain_dim = domain_range.shape[0]

        # check shape
        self.valid_shapes = ['box', 'annulus', 'circle', 'L-shape']
        if domain_shape in self.valid_shapes:
            pass
        else:
            raise ValueError(f'shape must be in {self.valid_shapes}')

        # different shape only valid in 2d
        if self.domain_dim == 2 or domain_shape == 'box':
            self.domain_shape = domain_shape
        else:
            raise NotImplementedError(f'{domain_shape} not implemented in {self.domain_dim}d')

        if mesh_domain is not None:
            x_train_domain = self.sample_domain_uniform(mesh_size=mesh_domain)
            self.x_train_domain = x_train_domain
        else:
            self.x_train_domain = None

        if mesh_boundary is not None:
            x_train_bc = self.sample_boundary_uniform(sample_size=mesh_boundary)
            self.x_train_bc = x_train_bc
        else:
            self.x_train_bc = None

    def __repr__(self):
        text = f'{self.domain_dim}d {self.domain_shape} domain with range: '
        for i in range(self.domain_dim):
            text += f'{self.domain_range[i]}' if i == 0 else f'*{self.domain_range[i]}'
        return text

    def sample_domain_uniform(self, mesh_size):
        """
           mesh_size: List
           若为int  如128   转化为128, 128 (维度为2的情况)
        """
        domain_dim = self.domain_dim

        if type(mesh_size) == int:
            mesh_size = [mesh_size] * domain_dim
        elif type(mesh_size) == list and len(mesh_size) == self.domain_dim:
            pass
        else:
            raise ValueError(f'mesh_size must be list of length {self.domain_dim} or int')

        # generate samples on [-1,1]^n
        x_train_domain_standard = self.sample_nd_mesh(mesh_size)

        # 根据不同形状去掉不符合条件的点
        if self.domain_shape == 'circle':
            assert self.domain_dim == 2, 'Dim should be 2 for circle.'
            r = np.sqrt(np.sum(x_train_domain_standard ** 2, axis=1))
            index1 = r <= 1.0
            x_train_domain_standard = x_train_domain_standard[index1, :]
            print(x_train_domain_standard.shape)

        if self.domain_shape == 'annulus':
            assert self.domain_dim == 2, 'Dim should be 2 for annulus.'
            r = np.sqrt(np.sum(x_train_domain_standard ** 2, axis=1))
            index1 = r <= 1.0
            index2 = r >= 0.5
            x_train_domain_standard = x_train_domain_standard[index1 * index2, :]

        if self.domain_shape == 'L-shape':
            assert self.domain_dim == 2, 'Dim should be 2 for L-shape.'
            index1 = np.sum(x_train_domain_standard > 0, axis=1) <= 1
            x_train_domain_standard = x_train_domain_standard[index1, :]

        # shift and scale by domain range
        x_train_domain = self.shift2range(x_train_domain_standard)
        return x_train_domain

    def sample_boundary_uniform(self, sample_size):
        """
           mesh_size: int
        """
        if self.domain_shape == 'circle':
            assert self.domain_dim == 2, 'Dim should be 2 for circle.'
            x_train_bc = self.sample_circle_uniform(r=1, sample_size=sample_size)

        if self.domain_shape == 'annulus':
            assert self.domain_dim == 2, 'Dim should be 2 for circle.'
            inner_size = int(sample_size / 3)
            outer_size = sample_size - inner_size
            x_train_boundary_out = self.sample_circle_uniform(r=1, sample_size=outer_size)
            x_train_boundary_in = self.sample_circle_uniform(r=0.5, sample_size=inner_size)
            x_train_bc = np.concatenate([x_train_boundary_out, x_train_boundary_in], axis=0)

        if self.domain_shape == 'L-shape':
            assert self.domain_dim == 2, 'Dim should be 2 for L-shape.'
            sample_1_side = int(np.ceil(sample_size / np.exp2(self.domain_dim))) + 1
            mesh_size = [sample_1_side] * self.domain_dim
            x_train_bc = self.sample_nd_mesh_bc(mesh_size=mesh_size)
            # reshape to L-shape
            x_train_bc = x_train_bc - 1.0 * (x_train_bc == 1) * (np.min(x_train_bc, axis=1, keepdims=True) > 0)

        if self.domain_shape == 'box':
            sample_1_side = sample_size/self.domain_dim/2
            mesh_size = [int(np.ceil(sample_1_side**(1.0/(self.domain_dim-1))))] * self.domain_dim
            x_train_bc = self.sample_nd_mesh_bc(mesh_size=mesh_size)

        # shift and scale by domain range
        x_train_bc = self.shift2range(x_train_bc)
        return x_train_bc

    def shift2range(self, x):
        lb = self.domain_range[:, 0]
        ub = self.domain_range[:, 1]
        x = (x * 0.5 + 0.5) * (ub - lb) + lb
        return x

    @staticmethod
    def sample_circle_uniform(r, sample_size):
        theta_vec = np.linspace(0, 2 * np.pi, sample_size + 1)[:-1]
        x_bc = np.stack([np.cos(theta_vec), np.sin(theta_vec)], axis=1) * r
        return x_bc

    @staticmethod
    def sample_nd_mesh_bc(mesh_size):
        # generate uniform samples on boundary of [-1,1]^n
        # mesh_size: list of int mesh for each dim
        domain_dim = len(mesh_size)
        if domain_dim >= 2:
            sample_all = []
            for i in range(domain_dim):
                bd_samples = Domain.sample_nd_mesh(np.delete(mesh_size, i))
                side1 = np.insert(bd_samples, obj=i, values=1, axis=1)
                side2 = np.insert(bd_samples, obj=i, values=-1, axis=1)
                sample_all.append(side1)
                sample_all.append(side2)
            # generate samples on [-1,1]^n
            box_bc = np.concatenate(sample_all, axis=0)
        else:
            box_bc = np.array([[-1], [1]])
        return box_bc

    @staticmethod
    def sample_nd_mesh(mesh_size):
        # generate uniform samples on [-1,1]^n
        # mesh_size: list of int mesh for each dim
        domain_dim = len(mesh_size)
        x_domain_standard = []
        temp_ones = np.ones(mesh_size)
        for i in range(domain_dim):
            mesh_1d = np.linspace(-1, 1, mesh_size[i])
            dim_array = np.ones(domain_dim, int)
            dim_array[i] = -1
            mesh_nd = mesh_1d.reshape(dim_array) * temp_ones
            x_domain_standard.append(mesh_nd.flatten())
        x_domain_standard = np.stack(x_domain_standard, axis=1)
        return x_domain_standard


def sample_boundaries(domain_range, num_samples):
    """
    在四个边界上均匀采样
    :param x_min: 最小 x 坐标
    :param y_min: 最小 y 坐标
    :param x_max: 最大 x 坐标
    :param y_max: 最大 y 坐标
    :param num_samples: 每个边界上的采样点数
    :return: top, bottom, left, right 四个边界上的采样点
    """
    x_min, x_max, y_min, y_max = domain_range
    # 顶部边界采样
    top_x = np.linspace(x_min, x_max, num_samples)
    top_y = np.full(num_samples, y_max)
    top = np.column_stack((top_x, top_y))

    # 底部边界采样
    bottom_x = np.linspace(x_min, x_max, num_samples)
    bottom_y = np.full(num_samples, y_min)
    bottom = np.column_stack((bottom_x, bottom_y))

    # 左侧边界采样
    left_x = np.full(num_samples, x_min)
    left_y = np.linspace(y_min, y_max, num_samples)
    left = np.column_stack((left_x, left_y))

    # 右侧边界采样
    right_x = np.full(num_samples, x_max)
    right_y = np.linspace(y_min, y_max, num_samples)
    right = np.column_stack((right_x, right_y))

    return top, bottom, left, right